import numpy as np
from chaco.api import *
from traits.api import *
from traitsui.api import *
from chaco.tools.api import PanTool, ZoomTool
from enable.api import ComponentEditor
import time
from PyQt5.QtCore import QThread


class PEC(HasTraits):
    nx = Int(100)
    ny = Int(100)
    s = Float(.5)
    nd = Int(20)
    nt = Int(5)
    stepNumber = Int(400)
    parameter_view = View(
        Item('nx', label='x方向网格数量'),
        Item('ny', label='y方向网格数量'),
        Item('s', label='时间步长因子'),
        Item('nd', label='空间步长因子'),
        Item('nt', label='曲线更新步数'),
        Item('stepNumber', label='时间步数'),
    )

class PML(HasTraits):

    nx = Int(100)
    ny = Int(100)
    s = Float(.5)
    nd = Int(20)
    na = Int(10)
    nt = Int(5)
    stepNumber = Int(400)
    parameter_view = View(
        Item('nx', label='x方向网格数量'),
        Item('ny', label='y方向网格数量'),
        Item('s', label='时间步长因子'),
        Item('nd', label='空间步长因子'),
        Item('nt', label='曲线更新步数'),
        Item('stepNumber', label='时间步数'),
        Item('na', label='吸收层厚度'),
    )

class Form(HasTraits):
    np.seterr(divide='ignore')

    paraPEC = Instance(PEC, ())
    paraPML = Instance(PML, ())
    method = Enum('pec', 'pml')
    plot = Instance(Plot)
    progress = Str('0%')
    thread = QThread()

    def run(self):
        match self.method:
            case 'pec':
                self.pec()
            case 'pml':
                self.pml()
    def pec(self):
        nx = self.paraPEC.nx
        ny = self.paraPEC.ny
        s = self.paraPEC.s
        nd = self.paraPEC.nd
        nt = self.paraPEC.nt
        e = np.zeros((nx + 1,ny + 1))
        hx = 0
        hy = 0
        for n in range(self.paraPEC.stepNumber + 1):
            hx -= s * np.diff(e[1:-1, :], 1, 1)
            hy += s * np.diff(e[:, 1:-1], 1, 0)
            e[1:-1, 1:-1] += s * (np.diff(hy, 1, 0) - np.diff(hx, 1, 1))
            e[nx // 2, ny // 2] += np.sin(2 * np.pi * n * s / nd)
            if n % nt == 0:
                self.plot.data.update_data(
                    e=e,
                )
                self.plot.title = f'Step number: {n}/{self.paraPEC.stepNumber}'
                time.sleep(.1)
                self.progress = f'{n / self.paraPEC.stepNumber:<.1%}'

    def pml(self):
        pass
#todo: 有待完善
    def traits_init(self):
        self.thread.run = self.run
        self.plot = Plot(ArrayPlotData())
        self.plot.data.update_data(
            e=np.arange(25).reshape((5, 5))
        )
        self.plot.img_plot('e')

    traits_view = View(
        UItem(
            "plot",
            editor=ComponentEditor(),
            resizable=True,
        ),
        HGroup(
            Item('progress', label='计算进度'),
            spring,
        ),
        width=1200,
        height=900,
        resizable=True,
    )
    parameter_view = View(
        HGroup(
            Item('method', style='custom'),
            spring
        ),
        UItem('paraPEC', style='custom', visible_when='method=="pec"'),
        UItem('paraPML', style='custom', visible_when='method=="pml"'),
        resizable=True,
        buttons=OKCancelButtons,
    )

if __name__ == "__main__":
    form = Form()
    form.thread.run()

